
import time

import SuperPocketCube
import gym
import numpy as np
import os


class Env(gym.Env):

    pocketCube: SuperPocketCube.SuperPocketCube
    state: np.ndarray

    def __init__(self):

        self.reset()

        # 状态空间？？？
        self.observation_space = gym.spaces.Box(low=0, high=5, shape=self.state.shape)

        # 动作空间
        self.action_space = gym.spaces.Discrete(12)
        self.actions = [i for i in range(12)]

    def step(self, action):
        done = False

        # last_chaos_level = self.pocketCube.chaosLevel
        self.pocketCube.rotate_action(action)
        now_chaos_level = self.pocketCube.chaosLevel
        now_chaos_level2 = self.pocketCube.chaosLevel2

        # determin reward
        reward = now_chaos_level + now_chaos_level2 - 1  # todo

        if now_chaos_level == 0:
            done = True
            reward = 1000000

        return self.get_state(), reward, done, # now_chaos_level

    def get_state(self):
        return np.array(self.pocketCube.observeFaces())

    def reset(self):
        self.pocketCube = SuperPocketCube.SuperPocketCube()
        self.pocketCube.shuffle()
        self.state = self.get_state()
        return self.state

    def render(self):
        formatted_value = self.pocketCube.show_front_up_right_view()
        os.system("cls")
        print(formatted_value)


if __name__ == '__main__':
    env = Env()
    obs = env.state
    done = False
    step_num = 1

    lowest_chaos = 1

    start_time = time.time()

    while not done:
        action = env.action_space.sample()
        next_obs, reward, done, chaos = env.step(action)
        # print(f"step#:{step_num} reward:{reward:>10.3f} done:{done} chaos:{chaos}")
        step_num += 1
        # env.render()

        if chaos < lowest_chaos:
            lowest_chaos = chaos
            print(f"step#:{step_num} reward:{reward:>10.3f} done:{done} chaos:{chaos}")

    end_time = time.time()
    time_use = end_time - start_time
    minute, sec = divmod(time_use, 60)
    print(f"\n用时 {int(minute)}分{int(sec): >2}秒")
